import os.path as osp

import torch.backends.cudnn as cudnn
from timm.utils import AverageMeter
from torch.cuda import amp

import utils.parallel as parallel_old
from networks.Fuck import *
from networks.evaluate import evaluate_main
from utils.criterion import *
from utils.one_class_criterion import *
from utils.utils import *

torch_ver = torch.__version__[:3]


class NetModelWithOutGan:
    def name(self):
        return 'kd_seg_without_gan'

    def DataParallelModelProcess(self, model, ParallelModelType=1, is_eval='train', device='cuda'):
        if ParallelModelType == 1:
            parallel_model = parallel_old.DataParallelModel(model)
        elif ParallelModelType == 2:
            parallel_model = parallel_old.DataParallelModel(model)
        else:
            raise ValueError('ParallelModelType should be 1 or 2')
        if is_eval == 'eval':
            parallel_model.eval()
        elif is_eval == 'train':
            parallel_model.train()
        else:
            raise ValueError('is_eval should be eval or train')
        parallel_model.float()
        parallel_model.to(device)
        return parallel_model

    def DataParallelCriterionProcess(self, criterion, device='cuda'):
        criterion = parallel_old.my_DataParallelCriterion(criterion)
        criterion.cuda()
        return criterion

    def __init__(self, args):
        cudnn.enabled = True
        self.args = args
        device = args.device

        student_model = Fuck7Res18 if args.classes_num == 1 else Fuck8Res18
        student = student_model(args.classes_num, pretrained=True)
        student.load_state_dict(torch.load(osp.join(self.args.snapshot_dir, 'model.pth')))
        print_model_parm_nums(student, 'student_model')
        self.parallel_student = self.DataParallelModelProcess(student, 2, 'train', device)
        self.student = student

        teacher_model = Fuck7 if args.classes_num == 1 else Fuck8
        teacher = teacher_model(args.classes_num)
        teacher.load_state_dict(torch.load("/home/hjr/zze/work_dirs/graduation/shandong/fuck7/model.pth")['net'])
        # teacher.load_state_dict(torch.load("/home/peter/zze/codes/work_dirs/graduation/whu/fuck8/model.pth")['net'])
        print_model_parm_nums(teacher, 'teacher_model')
        self.parallel_teacher = self.DataParallelModelProcess(teacher, 2, 'eval', device)
        self.teacher = teacher

        self.G_solver = torch.optim.AdamW(self.student.parameters(), lr=7e-5)

        self.best_mean_IU = args.best_mean_IU

        if args.classes_num != 1:
            self.criterion = self.DataParallelCriterionProcess(CriterionDSN())
            self.criterion_pixel_wise = self.DataParallelCriterionProcess(CriterionPixelWise())
            self.criterion_pair_wise_for_interfeat = self.DataParallelCriterionProcess(
                CriterionPairWiseforWholeFeatAfterPool(scale=args.pool_scale, feat_ind=2))
        else:
            self.criterion = self.DataParallelCriterionProcess(CriterionDSN_OneClass())
            self.criterion_pixel_wise = self.DataParallelCriterionProcess(CriterionPixelWise_OneClass())
            self.criterion_pair_wise_for_interfeat = self.DataParallelCriterionProcess(
                CriterionPairWiseforWholeFeatAfterPool_OneClass(scale=args.pool_scale, feat_ind=2))

        self.mc_G_loss = 0.0
        self.pi_G_loss = 0.0
        self.pa_G_loss = 0.0

        cudnn.benchmark = True
        if not os.path.exists(args.snapshot_dir):
            os.makedirs(args.snapshot_dir)

    def set_input(self, data):
        images, labels, _, _ = data
        self.images = images.cuda()
        self.labels = labels.long().cuda()

    def lr_poly(self, base_lr, iter, max_iter, power):
        return base_lr * ((1 - float(iter) / max_iter) ** power)

    def adjust_learning_rate(self, base_lr, optimizer, i_iter):
        args = self.args
        lr = self.lr_poly(base_lr, i_iter, args.num_steps, args.power)
        optimizer.param_groups[0]['lr'] = lr
        return lr

    def student_backward(self):
        temp = self.criterion(self.preds_S, self.labels, is_target_scattered=False)
        self.mc_G_loss = temp.item()
        G_loss = temp
        if self.args.pi:
            temp = self.args.lambda_pi * self.criterion_pixel_wise(self.preds_S, self.preds_T, is_target_scattered=True)
            self.pi_G_loss = temp.item()
            G_loss += temp
        if self.args.pa:
            temp1 = self.criterion_pair_wise_for_interfeat(self.preds_S, self.preds_T, is_target_scattered=True)
            self.pa_G_loss = temp1.item()
            G_loss += self.args.lambda_pa * temp1
        G_loss.backward()
        self.G_loss = G_loss.item()

    def optimize_parameters(self):
        with torch.no_grad():
            self.preds_T = self.parallel_teacher.eval()(self.images, parallel=self.args.parallel)
        self.preds_S = self.parallel_student.train()(self.images, parallel=self.args.parallel)

        self.G_solver.zero_grad()
        self.student_backward()
        self.G_solver.step()

    def evalute_model(self, model, loader, gpu_id, input_size, num_classes, whole):
        mean_IU, IU_array = evaluate_main(model=model,
                                          loader=loader,
                                          gpu_id=gpu_id,
                                          input_size=input_size,
                                          num_classes=num_classes,
                                          whole=whole)
        return mean_IU, IU_array

    def print_info(self, epoch, step):
        logging.info(
            'step:{:5d} G_lr:{:.6f} G_loss:{:.5f}(mc:{:.5f} pixelwise:{:.5f} pairwise:{:.5f})'.format(
                step, self.G_solver.param_groups[-1]['lr'],
                self.G_loss, self.mc_G_loss, self.pi_G_loss, self.pa_G_loss))

    def __del__(self):
        pass

    def save_ckpt(self, epoch, step, mean_IU, IU_array):
        torch.save(self.student.state_dict(), osp.join(self.args.snapshot_dir, 'model.pth'))
